# encoding=utf-8

import os
import cv2
import numpy as np
from PIL import Image
from seg_cuda.CudaMatrixTools import MatBuilder
from seg_system.application_config import ApplicationConfig
from seg_system.vascular.service.VascularToolsForBatch.VascularEachProcessor.VascularForBatchBase import \
    VascularForBatchBase


class CNBDForBatch(VascularForBatchBase):
    """CNBDForBatch:
        - 本算法对于并行度的支持不佳, 主要尝试使用 c++/openMP/thread的方案进行优化
        50张图的情况下, 默认开启Intersections并行时， 四种情况统计如下：
            - batch, stride, thread = 3, 15, 2 ( 测试环境 )
                - use_cxx, one_by_one, time_cost:
                - False  , False     , 180.58850858500227
                - False  , True      , 95.86277513900131
                - True   , False     , 14.0595488809995
                - True   , True      , 95.04049996100002
            - 精度损失状态(C++和python)
                - True   , False     , 0.863912080000091
            - 总结，python的多线程问题，需要使用语言扩展, 处理较大的矩阵的时候，应该选择其他方案
            - 调用语言扩展的时候，注意合并需要传输的数据，减小传输带来的大量消耗

        note:
            - Intersections部分需要重写为c++
    """

    def process_with_python(self, big_matrix: np.ndarray, usable_matrix: np.ndarray, each_matrix_shape: tuple,
                            stride: int, **kwargs):
        """尽量保证和VascularTools/CNBD一致
        """
        one_by_one = kwargs.get("one_by_one", False)
        usable_shape = usable_matrix.shape

        CNBD_list = []
        if not one_by_one:
            refine_output = self.hilditch(big_matrix)
        else:
            refine_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1]), np.uint8)
        branch_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1], 3), np.uint8)
        seg_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1], 3), np.uint8)

        for i in range(usable_shape[0]):
            for j in range(usable_shape[1]):
                if usable_matrix[i][j] == 0:
                    continue

                h_start = i * (each_matrix_shape[0] + stride)
                w_start = j * (each_matrix_shape[1] + stride)

                # each_mat = img
                each_mat = big_matrix[h_start: h_start + each_matrix_shape[0],
                           w_start: w_start + each_matrix_shape[1]]

                if one_by_one:
                    # XHout = img11
                    XHout = self.hilditch(each_mat)
                    refine_output[h_start: h_start + each_matrix_shape[0],
                    w_start: w_start + each_matrix_shape[1]] = XHout
                else:
                    XHout = refine_output[h_start: h_start + each_matrix_shape[0],
                            w_start: w_start + each_matrix_shape[1]]

                # 这部分没改好
                img111 = self.getSkeletonIntersection(XHout, ApplicationConfig.SystemConfig.VASCULAR_USE_WORKER)
                CNBD_list.append(len(img111))
                each_mat_copy = each_mat.copy()
                for item in img111:
                    data = np.array(item)
                    cv2.circle(each_mat_copy, (data[0], data[1]), 2, (0, 0, 255), 3)
                    cv2.circle(each_mat, (data[0], data[1]), 2, (0, 0, 0), 6)  # 用圆把节点断开

                branch_output[h_start: h_start + each_matrix_shape[0],
                w_start: w_start + each_matrix_shape[1]] = each_mat_copy
                seg_output[h_start: h_start + each_matrix_shape[0],
                w_start: w_start + each_matrix_shape[1]] = each_mat

        return MatBuilder.split(refine_output, usable_matrix, each_matrix_shape, stride), \
               MatBuilder.split(branch_output, usable_matrix, each_matrix_shape, stride), \
               MatBuilder.split(seg_output, usable_matrix, each_matrix_shape, stride), CNBD_list

    def process_with_cxx(self, big_matrix: np.ndarray, usable_matrix: np.ndarray, each_matrix_shape: tuple, stride: int,
                         **kwargs):
        one_by_one = kwargs.get("one_by_one", False)
        usable_shape = usable_matrix.shape

        CNBD_list = []
        if not one_by_one:
            refine_output = self.cxxPyVascular.hilditch(big_matrix)
        else:
            refine_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1]), np.uint8)
        branch_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1], 3), np.uint8)
        seg_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1], 3), np.uint8)

        for i in range(usable_shape[0]):
            for j in range(usable_shape[1]):
                if usable_matrix[i][j] == 0:
                    continue

                h_start = i * (each_matrix_shape[0] + stride)
                w_start = j * (each_matrix_shape[1] + stride)

                # each_mat = img
                each_mat = big_matrix[h_start: h_start + each_matrix_shape[0],
                           w_start: w_start + each_matrix_shape[1]]

                if one_by_one:
                    # XHout = img11
                    XHout = self.hilditch(each_mat)
                    refine_output[h_start: h_start + each_matrix_shape[0],
                    w_start: w_start + each_matrix_shape[1]] = XHout
                else:
                    XHout = refine_output[h_start: h_start + each_matrix_shape[0],
                            w_start: w_start + each_matrix_shape[1]]

                # 这部分没改好, C++和python有精度差,不知道为什么
                # img111_ = self.getSkeletonIntersection(XHout, ApplicationConfig.SystemConfig.VASCULAR_USE_WORKER)
                img111 = self.cxxPyVascular.GetSkeletonIntersection(XHout.copy())
                CNBD_list.append(len(img111))
                each_mat_copy = each_mat.copy()
                for item in img111:
                    data = np.array(item)
                    cv2.circle(each_mat_copy, (data[0], data[1]), 2, (0, 0, 255), 3)
                    cv2.circle(each_mat, (data[0], data[1]), 2, (0, 0, 0), 6)  # 用圆把节点断开

                branch_output[h_start: h_start + each_matrix_shape[0],
                w_start: w_start + each_matrix_shape[1]] = each_mat_copy
                seg_output[h_start: h_start + each_matrix_shape[0],
                w_start: w_start + each_matrix_shape[1]] = each_mat

        return MatBuilder.split(refine_output, usable_matrix, each_matrix_shape, stride), \
               MatBuilder.split(branch_output, usable_matrix, each_matrix_shape, stride), \
               MatBuilder.split(seg_output, usable_matrix, each_matrix_shape, stride), CNBD_list

    def batch_save(self, file_name: list, save_path: str, process_output, **kwargs):
        refine_output, branch_output, seg_output, CNBD_list = process_output

        refine_save_path = kwargs.get("refine_save_path")
        branch_save_path = kwargs.get("branch_save_path")
        seg_save_path = kwargs.get("seg_save_path")

        for e_n, e_r, e_b, e_s in zip(file_name, refine_output, branch_output, seg_output):
            e_r_p = os.path.join(refine_save_path, e_n)
            e_b_p = os.path.join(branch_save_path, e_n)
            e_s_p = os.path.join(seg_save_path, e_n)

            cv2.imwrite(e_r_p, e_r)
            cv2.imwrite(e_b_p, e_b)
            cv2.imwrite(e_s_p, e_s)

        return CNBD_list


class CNBDNoneBorderForBatch(CNBDForBatch):
    def process_with_python(self, big_matrix: np.ndarray, usable_matrix: np.ndarray, each_matrix_shape: tuple,
                            stride: int, **kwargs):
        one_by_one = kwargs.get("one_by_one", False)
        usable_shape = usable_matrix.shape

        CNBD_list = []
        if not one_by_one:
            refine_output = self.hilditch(big_matrix)
        else:
            refine_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1]), np.uint8)
        branch_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1], 3), np.uint8)
        seg_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1], 3), np.uint8)

        for i in range(usable_shape[0]):
            for j in range(usable_shape[1]):
                if usable_matrix[i][j] == 0:
                    continue

                h_start = i * (each_matrix_shape[0] + stride)
                w_start = j * (each_matrix_shape[1] + stride)

                # each_mat = img
                each_mat = big_matrix[h_start: h_start + each_matrix_shape[0],
                           w_start: w_start + each_matrix_shape[1]]

                if one_by_one:
                    # XHout = img11
                    XHout = self.hilditch(each_mat)
                    refine_output[h_start: h_start + each_matrix_shape[0],
                    w_start: w_start + each_matrix_shape[1]] = XHout
                else:
                    XHout = refine_output[h_start: h_start + each_matrix_shape[0],
                            w_start: w_start + each_matrix_shape[1]]

                # 2022.9.13新增的去除边框部分
                d = ApplicationConfig.SystemConfig.VASCULAR_DELETE_STRIDE
                XHout_mask = np.ones(XHout.shape, np.bool)
                XHout_mask[d: XHout.shape[0] - d, d: XHout_mask.shape[1] - d] = 0
                XHout[XHout_mask] = 0

                img111 = self.getSkeletonIntersection(XHout, ApplicationConfig.SystemConfig.VASCULAR_USE_WORKER)
                CNBD_list.append(len(img111))
                each_mat_copy = each_mat.copy()
                for item in img111:
                    data = np.array(item)
                    cv2.circle(each_mat_copy, (data[0], data[1]), 2, (0, 0, 255), 3)
                    cv2.circle(each_mat, (data[0], data[1]), 2, (0, 0, 0), 6)  # 用圆把节点断开

                branch_output[h_start: h_start + each_matrix_shape[0],
                w_start: w_start + each_matrix_shape[1]] = each_mat_copy
                seg_output[h_start: h_start + each_matrix_shape[0],
                w_start: w_start + each_matrix_shape[1]] = each_mat

        return MatBuilder.split(refine_output, usable_matrix, each_matrix_shape, stride), \
               MatBuilder.split(branch_output, usable_matrix, each_matrix_shape, stride), \
               MatBuilder.split(seg_output, usable_matrix, each_matrix_shape, stride), CNBD_list

    def process_with_cxx(self, big_matrix: np.ndarray, usable_matrix: np.ndarray, each_matrix_shape: tuple, stride: int,
                         **kwargs):
        one_by_one = kwargs.get("one_by_one", False)
        usable_shape = usable_matrix.shape

        CNBD_list = []
        if not one_by_one:
            refine_output = self.cxxPyVascular.hilditch(big_matrix)
        else:
            refine_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1]), np.uint8)
        branch_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1], 3), np.uint8)
        seg_output = np.zeros((big_matrix.shape[0], big_matrix.shape[1], 3), np.uint8)

        for i in range(usable_shape[0]):
            for j in range(usable_shape[1]):
                if usable_matrix[i][j] == 0:
                    continue

                h_start = i * (each_matrix_shape[0] + stride)
                w_start = j * (each_matrix_shape[1] + stride)

                # each_mat = img
                each_mat = big_matrix[h_start: h_start + each_matrix_shape[0],
                           w_start: w_start + each_matrix_shape[1]]

                if one_by_one:
                    # XHout = img11
                    XHout = self.hilditch(each_mat)
                    refine_output[h_start: h_start + each_matrix_shape[0],
                    w_start: w_start + each_matrix_shape[1]] = XHout
                else:
                    XHout = refine_output[h_start: h_start + each_matrix_shape[0],
                            w_start: w_start + each_matrix_shape[1]]

                # 2022.9.13新增的去除边框部分
                d = ApplicationConfig.SystemConfig.VASCULAR_DELETE_STRIDE
                XHout_mask = np.ones(XHout.shape, np.bool)
                XHout_mask[d: XHout.shape[0] - d, d: XHout_mask.shape[1] - d] = 0
                XHout[XHout_mask] = 0

                img111 = self.cxxPyVascular.GetSkeletonIntersection(XHout.copy())
                CNBD_list.append(len(img111))
                each_mat_copy = each_mat.copy()
                for item in img111:
                    data = np.array(item)
                    cv2.circle(each_mat_copy, (data[0], data[1]), 2, (0, 0, 255), 3)
                    cv2.circle(each_mat, (data[0], data[1]), 2, (0, 0, 0), 6)  # 用圆把节点断开

                branch_output[h_start: h_start + each_matrix_shape[0],
                w_start: w_start + each_matrix_shape[1]] = each_mat_copy
                seg_output[h_start: h_start + each_matrix_shape[0],
                w_start: w_start + each_matrix_shape[1]] = each_mat

        return MatBuilder.split(refine_output, usable_matrix, each_matrix_shape, stride), \
               MatBuilder.split(branch_output, usable_matrix, each_matrix_shape, stride), \
               MatBuilder.split(seg_output, usable_matrix, each_matrix_shape, stride), CNBD_list

    def batch_save(self, file_name: list, save_path: str, process_output, **kwargs):
        return super().batch_save(file_name, save_path, process_output, **kwargs)
